import pandas as pd
import numpy as np
import rasterio
from sklearn.ensemble import RandomForestRegressor
from sklearn.preprocessing import StandardScaler
from sklearn.model_selection import train_test_split
import warnings
import os
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
from matplotlib import cm

# 忽略警告
warnings.filterwarnings("ignore")

# 设置工作目录
data_dir = r"D:\data1"

# 读取公益林数据（仅提取必要的列）
forest_data = pd.read_excel(
    os.path.join(data_dir, "2021公益林特征数据（2024.12）.xls"),
    usecols=["蓄积量", "横坐标", "纵坐标"]
)

# 提取2014年NDVI栅格数据（仅用于模型训练）
with rasterio.open(os.path.join(data_dir, "hunan_ndvi_2014.tif")) as src:
    ndvi_profile = src.profile
    ndvi_2014_data = src.read(1)

# 定义函数提取NDVI值
def extract_ndvi(row, ndvi_data, src):
    try:
        lon, lat = row['纵坐标'], row['横坐标']
        row_idx, col_idx = src.index(lon, lat)
        if 0 <= row_idx < ndvi_data.shape[0] and 0 <= col_idx < ndvi_data.shape[1]:
            return ndvi_data[row_idx, col_idx]
        else:
            return np.nan
    except:
        return np.nan

# 为训练数据提取NDVI值
with rasterio.open(os.path.join(data_dir, "hunan_ndvi_2014.tif")) as src:
    forest_data['extracted_ndvi'] = forest_data.apply(
        lambda row: extract_ndvi(row, ndvi_2014_data, src), axis=1
    )

# 数据清洗
forest_data = forest_data.dropna(subset=['extracted_ndvi', '蓄积量'])

# 特征选择与目标变量
features = ['extracted_ndvi']
X = forest_data[features]
y = forest_data['蓄积量']

# 划分训练集和测试集
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

# 特征标准化
scaler = StandardScaler()
X_train_scaled = scaler.fit_transform(X_train)
X_test_scaled = scaler.transform(X_test)

# 训练优化后的随机森林模型
model = RandomForestRegressor(
    n_estimators=50,  # 减少树的数量
    max_depth=8,      # 限制树的深度
    min_samples_split=10,
    min_samples_leaf=3,
    random_state=42,
    n_jobs=-1         # 使用所有CPU核心
)
model.fit(X_train_scaled, y_train)

# 定义函数处理单一年份的预测
def process_year(year):
    ndvi_file = os.path.join(data_dir, f"hunan_ndvi_{year}.tif")
    with rasterio.open(ndvi_file) as src:
        ndvi_data = src.read(1)
        
    # 创建预测特征矩阵
    ndvi_flat = ndvi_data.flatten()
    valid_mask = ~np.isnan(ndvi_flat)
    
    # 创建预测数据
    X_grid = np.zeros((len(ndvi_flat), len(features)))
    X_grid[:, 0] = ndvi_flat
    X_grid[~valid_mask] = 0  # 填充缺失值
    
    # 标准化预测数据
    X_grid_scaled = scaler.transform(X_grid)
    
    # 预测生物量
    biomass_pred = model.predict(X_grid_scaled)
    biomass_pred = biomass_pred.reshape(ndvi_data.shape)
    return biomass_pred

# 使用并行处理加速多年预测
years = np.arange(2014, 2025)
biomass_data = [process_year(year) for year in years]

# 获取栅格数据的地理范围
with rasterio.open(os.path.join(data_dir, "hunan_ndvi_2014.tif")) as src:
    bounds = src.bounds
lon = np.linspace(bounds.left, bounds.right, num=biomass_data[0].shape[1])
lat = np.linspace(bounds.bottom, bounds.top, num=biomass_data[0].shape[0])

# 3D可视化优化版
plt.style.use('classic')  # 使用经典风格
fig = plt.figure(figsize=(14, 10))
ax = fig.add_subplot(111, projection='3d')

# 使用降采样加速可视化
sample_step = 50
lon_sampled = lon[::sample_step]
lat_sampled = lat[::sample_step]

# 更优的绘图选项
for i, year in enumerate(years):
    z = biomass_data[i][::sample_step, ::sample_step]
    lon_grid, lat_grid = np.meshgrid(lon_sampled, lat_sampled)
    # 使用更丰富的颜色映射
    surf = ax.plot_surface(
        lon_grid, lat_grid, z, rstride=1, cstride=1, alpha=0.25,
        cmap='viridis', edgecolor='none', linewidth=0.2
    )

# 添加颜色条
fig.colorbar(surf, ax=ax, shrink=0.5, aspect=10, label='Biomass Value')

# 设置更美观的标签和标题
ax.set_xlabel('Longitude', fontsize=12, fontfamily='sans-serif')
ax.set_ylabel('Latitude', fontsize=12, fontfamily='sans-serif')
ax.set_zlabel('Biomass Value', fontsize=12, fontfamily='sans-serif')

# 添加网格线
ax.grid(True, which='both', axis='y', linestyle='--', linewidth=0.5)
ax.xaxis._axinfo['grid'].update(visible=False)

# 调整视角
ax.view_init(elev=30, azim=120)

# 设置背景透明
ax.set_facecolor((0, 0, 0, 0))

output_path = os.path.join(data_dir, "optimized_hunan_biomass_3d.png")
plt.savefig(output_path, dpi=300, bbox_inches='tight', transparent=True)
plt.show()
print(f"优化后的3D可视化结果已保存至: {output_path}")